import os

from PIL import Image
from tqdm import tqdm

from option import parse_args

MAX_ACTOR = 5


def cmp_video_name(video_name):
    video_name_splt = video_name.split('_')
    if len(video_name_splt) == 1:
        return int(video_name[1])
    elif len(video_name_splt) == 2:
        return int(video_name_splt[0][2:]) + 300
    else:
        return int(video_name_splt[0][2:]) * 1000 + int(video_name_splt[1][2:]) * 10 + int(video_name_splt[2])


def count_actors(video_path, frame):
    num = 0
    for i in range(MAX_ACTOR):
        temp_path = os.path.join(video_path, "{}_{}.png".format(frame, i))
        if os.path.exists(temp_path):
            num += 1
    return num


def get_frame_min_max_width(video_path, frame):
    w_min, w_max = 500, 0
    for i in range(MAX_ACTOR):
        image_path = os.path.join(video_path, "{}_{}.png".format(frame, i))
        if os.path.exists(image_path):
            with Image.open(image_path) as img:
                w_min = min(w_min, img.size[0])
                w_max = max(w_max, img.size[0])
    return w_min, w_max


def multiple_actors(crops, func, template):
    video_names = os.listdir(crops)
    video_names.sort(key=cmp_video_name)
    for video_name in tqdm(video_names):
        video_path = os.path.join(crops, video_name)
        num = 1
        prev_frame = '-1'
        width_pairs = []
        for crop_name in os.listdir(video_path):
            crop_name_split = crop_name[:-4].split('_')
            assert len(crop_name_split) == 2
            if int(crop_name_split[1]) == 0 or crop_name_split[0] == prev_frame:
                continue
            actors = count_actors(video_path, crop_name_split[0])
            num = max(num, actors)
            prev_frame = crop_name_split[0]
            w_min, w_max = get_frame_min_max_width(video_path, crop_name_split[0])
            width_pairs.append((w_min, w_max))
        if num > 1:
            func(template.format(video_name, num))


def main():
    args = parse_args()
    crops = os.path.join(args.root_dir, 'crops')
    args.save = True
    if args.save:
        with open('multiple_actors_statistics.dat', 'w') as f:
            multiple_actors(crops, f.write, '{},{}\n')
    else:
        multiple_actors(crops, print, '{},{}')


if __name__ == '__main__':
    main()
